import torch
from ResNeSt import ResNeSt
from ResNestBottleneck import ResNestBottleneck
from train import train_model


def resnest50(num_classes=100, **kwargs):
    return ResNeSt(ResNestBottleneck, [3, 4, 6, 3], num_classes=num_classes,
                  radix=2, cardinality=2, base_width=64, **kwargs)

def resnest18(num_classes=100, **kwargs):
    return ResNeSt(ResNestBottleneck, [2, 2, 2, 2], num_classes=num_classes,
                  radix=2, cardinality=1, base_width=64, **kwargs)


if __name__ == "__main__":
    # 创建 ResNeSt 模型
    
    model = resnest50(num_classes=100)
    print(f"ResNeSt-50 模型参数量: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")

    # 训练模型
    model = train_model(model, epochs=15, lr=0.001)
    
    # 保存最终模型
    torch.save({
        'model_state_dict': model.state_dict(),
    }, 'resnest_final.pth')
    
    print("训练完成!")